import torch

from problems.DTLZ import DTLZ


class DTLZ1(DTLZ):
    def __init__(self, var_dim: int, obj_dim: int, max_fun_eval, kwargs: dict):
        super().__init__(var_dim, obj_dim, max_fun_eval, kwargs)

    def eval_value(self, dec):
        g = super().get_g(dec)
        obj = 0.5 * g * torch.fliplr(
            torch.cumprod(torch.cat((torch.ones(dec.size(0), 1), dec[:, :self.obj_dim - 1]), dim=1),
                          dim=1)) * torch.cat(
            (torch.ones(dec.size(0), 1), 1 - torch.fliplr(dec[:, :self.obj_dim - 1])), dim=1)
        return obj

    def get_optimal_solutions(self, size):
        return super().get_optimal_solutions0(size)
